package cz.cuni.mff.ufal.volk.langr;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.Reader;
import java.io.Serializable;
import java.security.InvalidParameterException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Vector;
import java.util.Map.Entry;

/**
 * <p>An n-gram language model.</p>
 * <p>The probabilities of particular n-grams are updated during the calls of {@link #addString(String)} 
 * and {@link #addReader(Reader)} methods.</p>
 * <p>The probabilities are counted after the sequence of the n-gram count updates. To do it, you should call 
 * the {@link #updateProbabilities()} method.</p>
 * <p>It is advised to estimate the smoothing parameters. To do it, you should call the {@link #smoothen(Map, double, int)}
 * method.</p>
 * 
 * @author Bartłomiej Etenkowski
 */
public class NgramLanguageModel implements Serializable {   
  private static final long serialVersionUID = 3807753089334327047L;

  /**
   * Constructs the language model of the given order.
   * @param languageId the language ID
   * @param n the order of the language model
   * @throws NullPointerException if the <tt>languageId</tt> parameter is null
   * @throws InvalidParameterException if the <tt>@link #n</tt> parameter is less than <tt>1</tt>
   */
  public NgramLanguageModel(String languageId, int n) {
    languageId.getClass();
    if (n < 1)
      throw new InvalidParameterException("The n parameter is less than 1");
    
    this.languageId = languageId;
    this.n = n;
    this.ngramCounts = new Vector<Map<String, Long>>();
    for (int i = 0; i <= n; i++)
      this.ngramCounts.add(new HashMap<String, Long>());
    this.frequencies = new Vector<Map<String,Double>>();
    for (int i = 0; i <= n; i++)
      this.frequencies.add(new HashMap<String, Double>());
    this.lambda = new double[n + 1];
    for (int i = 0; i <= n; i++)
      this.lambda[i] = 1. / (n + 1);
  }

  private final String languageId;

  /**
   * @return the language id
   */
  public String getLanguageId() {
    return languageId;
  }
  
  /** the model order */
  private final int n;
  
  /**
   * Returns the model order.
   * @return the model order
   */
  public int getModelOrder() {
    return n;
  }
  
  /** 
   * The n-gram counts. The element <tt>i</tt> contains i-gram counts. The map within the element 0 is always
   * empty.
   */
  private final Vector<Map<String, Long>> ngramCounts;
  
  /** P(0) */
  private double prob0;
 
  /**
   * <p>The n-gram frequencies (raw probabilities).</p> 
   */
  private final Vector<Map<String, Double>> frequencies;
  
  public Map<String, Double> getFrequencies(int n) {
    return Collections.unmodifiableMap(frequencies.elementAt(n));
  }
  
  public Map<Character, Double> getCharacterProbabilities() {
    if (n < 1)
      throw new UnsupportedOperationException();
    Map<Character, Double> probs = new HashMap<Character, Double>(frequencies.elementAt(1).size());
    for (Entry<String, Double> e : frequencies.elementAt(1).entrySet())
      probs.put(e.getKey().charAt(0), e.getValue());
    return probs;
  }
  
  /**
   * Updates the n-gram counts - adds the n-grams contained by the parameter <tt>s</tt>.
   * @param s
   */
  public void addString(String s) {
    for (int i = 0; i < s.length(); i++) {
      for (int j = 1; j <= n && i + j <= s.length(); j++) {
        String key = s.substring(i, i + j);
        Utils.incrementLongValue(ngramCounts.elementAt(j), key);
      }
    }
  }
  
  /**
   * Reads the lines from the {@link Reader} <tt>r</tt> while available and updates the n-gram counts. 
   * The method does <b>not</b> close the reader when finished. 
   * @param r the reader
   * @throws IOException if an {@code IOException} occurred
   */
  public void addReader(Reader r) throws IOException {
    BufferedReader br = new BufferedReader(r);
    String line;
    while ((line = br.readLine()) != null)
      addString(line);
  }
  
  /**
   * Updates the n-gram probabilities after updating the n-gram counts. The method should be called after
   * the sequence of the {@link #addString(String)} and {@link #addReader(Reader)} calls.0 
   */
  public void updateProbabilities() {
    for (int i = 1; i <= n; i++) {
      Map<String, Long> counts = ngramCounts.elementAt(i);
      Map<String, Double> probs = frequencies.elementAt(i);
      probs.clear();
      for (Entry<String, Long> e : counts.entrySet()) {
        double dcount = (double)e.getValue();
        probs.put(e.getKey(), dcount);
      }
      if (i == 1) {
        long sum = 0L;
        for (long l : counts.values())
          sum += l;
        for (Entry<String, Double> e : probs.entrySet()) {
          e.setValue(e.getValue() / sum);
        }
      } else {
        for (Entry<String, Double> e : probs.entrySet()) {
          String ngram = e.getKey();
          e.setValue(e.getValue() / ngramCounts.elementAt(i - 1).get(ngram.substring(0, ngram.length() - 1)));
        }
      }
    }
    prob0 = 1. / frequencies.elementAt(1).size();
    frequencies.elementAt(0).put("", prob0);
  }
  
  /** the smoothing parameters */
  private double[] lambda;
  
  /**
   * Retrieves the raw (not smoothed) probability of the given n-gram (P(C|Hist)).
   * @param text the text which contains the n-gram
   * @param beginIndex the begin index of the n-gram in the text
   * @param order the order of the n-gram
   * @return the raw probability (without smoothing) of the n-gram
   * @throws InvalidParameterException if the <tt>length</tt> parameter is not between 0 and n
   */
  public double getRawProbability(String text, int beginIndex, int order) {
    if (order < 0 || order > n)
      throw new InvalidParameterException("The length parameter is not between 0 and n");
    if (order == 0) {
      return prob0;
    } else {
      Double prob = frequencies.elementAt(order).get(text.substring(beginIndex, beginIndex + order));
      if (prob == null)
        return 0.;
      else
        return prob;
    }
  }
  
  /**
   * Retrieves the smoothed probability of the given n-gram.
    *@param text the text which contains the n-gram
   * @param beginIndex the begin index of the n-gram in the text
   * @return the smoothed probability of the n-gram
   */
  public double getSmoothedProbability(String text, int beginIndex) {
    if (beginIndex > text.length() + n)
      return 0.;
    double prob = 0.;
    for (int i = 0; i <= n; i++)
      prob += lambda[i] * getRawProbability(text, beginIndex + n - i, i);
    return prob;
  }
  
  private double getSmoothedProbability(String text, int beginIndex, int length) {
    if (length == n) {
      return getSmoothedProbability(text, beginIndex);
    } else {
      if (beginIndex + length > text.length())
        return 0.;
      double[] lambda = new double[length + 1];
      double sum = 0.;
      for (int i = 0; i <= length; i++) {
        lambda[i] = this.lambda[i];
        sum += this.lambda[i];
      }
      for (int i = 0; i <= length; i++) {
        if (sum > 0) {
          lambda[i] /= sum;
        } else {
          lambda[i] = 1. / (length + 1);
        }
      }
      double prob =0.;
      for (int i = 0; i <= length; i++)
        prob += lambda[i] * getRawProbability(text, beginIndex + length - i, i);
      return prob;
    }
  }
  
  /**
   * Retrieves the smoothed probability of the n-gram. The call <tt>getSmoothedProbability(ngram)</tt>
   * is equivalent to <tt>getSmoothedProbability(ngram, 0)</tt>.
   * @param ngram the n-gram
   * @return the smoothed probability of the n-gram in this language model
   */
  public double getSmoothedProbability(String ngram) {
    return getSmoothedProbability(ngram, 0);
  }
  
  /**
   * Returns the log probability of the given text.
   * @param text
   * @return the log probability of the given text
   */
  public double logProbability(String text) {
    if (text.length() == 0) {
      return Double.NEGATIVE_INFINITY;
    } else {
      double logProb = 0.;
      for (int i = 0; i < text.length(); i++) {
        int beginIndex = i - n < 0 ? 0 : i - n;
        int length = i >= n ? n : i;
        logProb += Math.log(getSmoothedProbability(text, beginIndex, length));
      }
      return logProb;
    }
  }
  
  /**
   * Estimation of the smoothing parameters (EM algorithm).
   * @param ngramCounts the map of the n-gram counts used to estimate the smoothing parameters. The counts must <b>not</b>
   *                    be estimated from the data used to count n-grams.
   * @param epsilon
   * @param maxIterations the maximal count of iterations
   * @throws NullPointerException if the <tt>ngramCounts</tt> parameter is <tt>null</tt>
   * @throws InvalidParameterException if the <tt>ngramCounts</tt> map contains one or more n-grams of wrong order 
   */
  public void smoothen(Map<String, Integer> ngramCounts, double epsilon, int maxIterations) {
    ngramCounts.getClass(); // null check
    if (epsilon < 0. || epsilon > 1.)
      throw new InvalidParameterException("The epsilon parameter must be between 0. and 1.");
    if (maxIterations < 1)
      throw new InvalidParameterException("The maxIterations parameter is less than 1");
    
    double delta = 1.;
    for (int i = 0; i < maxIterations && delta > epsilon; i++) {
      System.out.printf("smoothing: %d\t", i + 1);
      for (int j = 0; j <= n; j++)
        System.out.printf("l[%d]=%f ", j, lambda[j]);
      System.out.println();
      
      double[] c = new double[n + 1]; // new lambda candidates
      for (Entry<String, Integer> e : ngramCounts.entrySet()) {
        String ngram = e.getKey();
        if (ngram == null || ngram.length() != n)
          throw new InvalidParameterException("The ngramCounts map contains wrong n-grams");
        if (e.getValue() == null || e.getValue() < 0)
          throw new InvalidParameterException("The ngramCounts map contains null or negative counts");
        int count = e.getValue();
        double smProb = getSmoothedProbability(ngram);
        for (int j = 0; j <= n; j++) {
          c[j] += lambda[j] * count * getRawProbability(ngram, n - j, j) / smProb;
        }
      }
      
      delta = 0.;
      double sum = 0.;
      for (double d : c)
        sum += d;
      if (sum <= 0) {
        System.out.printf("sum = %f!\n", sum);
        break;
      }
      for (int j = 0; j <= n; j++) {
        double nval = c[j] / sum;
        double diff = Math.abs(nval - lambda[j]);
        if (diff > delta)
          delta = diff;
        lambda[j] = nval;
      }
    }
  }
  
  @Override
  public String toString() {
    return String.format("[%s] - LM", languageId);
  }
}
